// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

//! 文件格式和序列化支持

use std::fs::{File, OpenOptions};
use std::io::{Read, Write, Seek, SeekFrom, BufReader, BufWriter};
use std::path::Path;
use bytes::{Bytes, BytesMut, BufMut};

use crate::ipc::recorder::proto_types::{
    RecordFileHeader, RecordBatch, RecordFileFooter, RequestRecord,
    SystemContext, RecordStatistics
};
use super::{RecorderError, RecordFormat};

/// 记录文件魔数：0x52455152 ("REQR")
pub const RECORD_FILE_MAGIC: u32 = 0x52455152;

/// 记录文件版本
pub const RECORD_FILE_VERSION: &str = "1.0";

/// 记录写入器trait
pub trait RecordWriter {
    /// 写入文件头部
    fn write_header(&mut self, header: &RecordFileHeader) -> Result<(), RecorderError>;
    
    /// 写入记录批次
    fn write_batch(&mut self, batch: &RecordBatch) -> Result<(), RecorderError>;
    
    /// 写入文件尾部
    fn write_footer(&mut self, footer: &RecordFileFooter) -> Result<(), RecorderError>;
    
    /// 刷新缓冲区
    fn flush(&mut self) -> Result<(), RecorderError>;
    
    /// 获取当前写入位置
    fn position(&mut self) -> Result<u64, RecorderError>;
    
    /// 同步到磁盘
    fn sync(&mut self) -> Result<(), RecorderError>;
}

/// 记录读取器trait
pub trait RecordReader {
    /// 读取文件头部
    fn read_header(&mut self) -> Result<RecordFileHeader, RecorderError>;
    
    /// 读取下一个记录批次
    fn read_batch(&mut self) -> Result<Option<RecordBatch>, RecorderError>;
    
    /// 读取文件尾部
    fn read_footer(&mut self) -> Result<RecordFileFooter, RecorderError>;
    
    /// 跳转到指定位置
    fn seek(&mut self, pos: u64) -> Result<u64, RecorderError>;
    
    /// 获取当前读取位置
    fn position(&mut self) -> Result<u64, RecorderError>;
    
    /// 获取文件大小
    fn file_size(&mut self) -> Result<u64, RecorderError>;
}

/// Protobuf格式写入器
pub struct ProtobufWriter {
    writer: BufWriter<File>,
}

impl ProtobufWriter {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, RecorderError> {
        let file = OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(path)?;
        
        Ok(Self {
            writer: BufWriter::new(file),
        })
    }
    
    /// 写入长度前缀的JSON消息
    fn write_message<M>(&mut self, msg: &M) -> Result<(), RecorderError> 
    where 
        M: serde::Serialize,
    {
        let json = serde_json::to_vec(msg)
            .map_err(|e| RecorderError::Serialization(e.to_string()))?;
        
        // 写入消息长度（4字节大端序）
        let len = json.len() as u32;
        self.writer.write_all(&len.to_be_bytes())?;
        
        // 写入消息内容
        self.writer.write_all(&json)?;
        
        Ok(())
    }
}

impl RecordWriter for ProtobufWriter {
    fn write_header(&mut self, header: &RecordFileHeader) -> Result<(), RecorderError> {
        self.write_message(header)
    }
    
    fn write_batch(&mut self, batch: &RecordBatch) -> Result<(), RecorderError> {
        self.write_message(batch)
    }
    
    fn write_footer(&mut self, footer: &RecordFileFooter) -> Result<(), RecorderError> {
        self.write_message(footer)
    }
    
    fn flush(&mut self) -> Result<(), RecorderError> {
        self.writer.flush().map_err(RecorderError::Io)
    }
    
    fn position(&mut self) -> Result<u64, RecorderError> {
        self.writer.stream_position().map_err(RecorderError::Io)
    }
    
    fn sync(&mut self) -> Result<(), RecorderError> {
        self.writer.get_mut().sync_all().map_err(RecorderError::Io)
    }
}

/// Protobuf格式读取器
pub struct ProtobufReader {
    reader: BufReader<File>,
}

impl ProtobufReader {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, RecorderError> {
        let file = File::open(path)?;
        
        Ok(Self {
            reader: BufReader::new(file),
        })
    }
    
    /// 读取长度前缀的JSON消息
    fn read_message<M>(&mut self) -> Result<Option<M>, RecorderError> 
    where
        M: serde::de::DeserializeOwned,
    {
        // 读取消息长度
        let mut len_buf = [0u8; 4];
        match self.reader.read_exact(&mut len_buf) {
            Ok(()) => {},
            Err(ref e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
            Err(e) => return Err(RecorderError::Io(e)),
        }
        
        let len = u32::from_be_bytes(len_buf) as usize;
        if len > 100 * 1024 * 1024 { // 100MB 限制
            return Err(RecorderError::FormatError { 
                message: format!("Message too large: {} bytes", len)
            });
        }
        
        // 读取消息内容
        let mut buf = vec![0u8; len];
        self.reader.read_exact(&mut buf)?;
        
        // 解析消息
        let msg: M = serde_json::from_slice(&buf)
            .map_err(|e| RecorderError::Deserialization(e.to_string()))?;
        Ok(Some(msg))
    }
}

impl RecordReader for ProtobufReader {
    fn read_header(&mut self) -> Result<RecordFileHeader, RecorderError> {
        self.read_message()?
            .ok_or_else(|| RecorderError::FormatError { 
                message: "Missing file header".to_string()
            })
    }
    
    fn read_batch(&mut self) -> Result<Option<RecordBatch>, RecorderError> {
        self.read_message()
    }
    
    fn read_footer(&mut self) -> Result<RecordFileFooter, RecorderError> {
        self.read_message()?
            .ok_or_else(|| RecorderError::FormatError { 
                message: "Missing file footer".to_string()
            })
    }
    
    fn seek(&mut self, pos: u64) -> Result<u64, RecorderError> {
        self.reader.seek(SeekFrom::Start(pos)).map_err(RecorderError::Io)
    }
    
    fn position(&mut self) -> Result<u64, RecorderError> {
        self.reader.stream_position().map_err(RecorderError::Io)
    }
    
    fn file_size(&mut self) -> Result<u64, RecorderError> {
        let current_pos = self.position()?;
        let size = self.reader.seek(SeekFrom::End(0))?;
        self.reader.seek(SeekFrom::Start(current_pos))?;
        Ok(size)
    }
}

/// JSON格式写入器
pub struct JsonWriter {
    writer: BufWriter<File>,
    first_record: bool,
}

impl JsonWriter {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, RecorderError> {
        let file = OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(path)?;
        
        let mut writer = BufWriter::new(file);
        writer.write_all(b"{\n")?; // 开始JSON对象
        
        Ok(Self {
            writer,
            first_record: true,
        })
    }
}

impl RecordWriter for JsonWriter {
    fn write_header(&mut self, header: &RecordFileHeader) -> Result<(), RecorderError> {
        let json = serde_json::to_string_pretty(header)
            .map_err(|e| RecorderError::Serialization(e.to_string()))?;
        
        write!(self.writer, "  \"header\": {},\n", json)?;
        write!(self.writer, "  \"records\": [\n")?;
        
        Ok(())
    }
    
    fn write_batch(&mut self, batch: &RecordBatch) -> Result<(), RecorderError> {
        for record in &batch.records {
            if !self.first_record {
                write!(self.writer, ",\n")?;
            }
            self.first_record = false;
            
            let json = serde_json::to_string_pretty(record)
                .map_err(|e| RecorderError::Serialization(e.to_string()))?;
            
            write!(self.writer, "    {}", json)?;
        }
        
        Ok(())
    }
    
    fn write_footer(&mut self, footer: &RecordFileFooter) -> Result<(), RecorderError> {
        write!(self.writer, "\n  ],\n")?;
        
        let json = serde_json::to_string_pretty(footer)
            .map_err(|e| RecorderError::Serialization(e.to_string()))?;
        
        write!(self.writer, "  \"footer\": {}\n", json)?;
        write!(self.writer, "}}\n")?; // 结束JSON对象
        
        Ok(())
    }
    
    fn flush(&mut self) -> Result<(), RecorderError> {
        self.writer.flush().map_err(RecorderError::Io)
    }
    
    fn position(&mut self) -> Result<u64, RecorderError> {
        self.writer.stream_position().map_err(RecorderError::Io)
    }
    
    fn sync(&mut self) -> Result<(), RecorderError> {
        self.writer.get_mut().sync_all().map_err(RecorderError::Io)
    }
}

/// JSON格式读取器
pub struct JsonReader {
    content: serde_json::Value,
    record_index: usize,
}

impl JsonReader {
    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, RecorderError> {
        let mut file = File::open(path)?;
        let mut content = String::new();
        file.read_to_string(&mut content)?;
        
        let json: serde_json::Value = serde_json::from_str(&content)
            .map_err(|e| RecorderError::Deserialization(e.to_string()))?;
        
        Ok(Self {
            content: json,
            record_index: 0,
        })
    }
}

impl RecordReader for JsonReader {
    fn read_header(&mut self) -> Result<RecordFileHeader, RecorderError> {
        let header_json = self.content.get("header")
            .ok_or_else(|| RecorderError::FormatError { 
                message: "Missing header in JSON file".to_string()
            })?;
        
        serde_json::from_value(header_json.clone())
            .map_err(|e| RecorderError::Deserialization(e.to_string()))
    }
    
    fn read_batch(&mut self) -> Result<Option<RecordBatch>, RecorderError> {
        let records = self.content.get("records")
            .and_then(|v| v.as_array())
            .ok_or_else(|| RecorderError::FormatError { 
                message: "Missing or invalid records array".to_string()
            })?;
        
        if self.record_index >= records.len() {
            return Ok(None);
        }
        
        // 读取单个记录作为批次
        let record_json = &records[self.record_index];
        self.record_index += 1;
        
        let record: RequestRecord = serde_json::from_value(record_json.clone())
            .map_err(|e| RecorderError::Deserialization(e.to_string()))?;
        
        let batch = RecordBatch {
            records: vec![record],
            batch_timestamp_ns: 0,
            batch_id: self.record_index as u32,
        };
        
        Ok(Some(batch))
    }
    
    fn read_footer(&mut self) -> Result<RecordFileFooter, RecorderError> {
        let footer_json = self.content.get("footer")
            .ok_or_else(|| RecorderError::FormatError { 
                message: "Missing footer in JSON file".to_string()
            })?;
        
        serde_json::from_value(footer_json.clone())
            .map_err(|e| RecorderError::Deserialization(e.to_string()))
    }
    
    fn seek(&mut self, _pos: u64) -> Result<u64, RecorderError> {
        // JSON格式不支持随机访问
        Err(RecorderError::FormatError { 
            message: "Seek not supported for JSON format".to_string()
        })
    }
    
    fn position(&mut self) -> Result<u64, RecorderError> {
        Ok(self.record_index as u64)
    }
    
    fn file_size(&mut self) -> Result<u64, RecorderError> {
        let records = self.content.get("records")
            .and_then(|v| v.as_array())
            .ok_or_else(|| RecorderError::FormatError { 
                message: "Missing or invalid records array".to_string()
            })?;
        
        Ok(records.len() as u64)
    }
}

/// 根据格式创建记录写入器
pub fn create_writer<P: AsRef<Path>>(path: P, format: RecordFormat) -> Result<Box<dyn RecordWriter>, RecorderError> {
    match format {
        RecordFormat::Protobuf => Ok(Box::new(ProtobufWriter::new(path)?)),
        RecordFormat::Json => Ok(Box::new(JsonWriter::new(path)?)),
    }
}

/// 根据格式创建记录读取器
pub fn create_reader<P: AsRef<Path>>(path: P, format: RecordFormat) -> Result<Box<dyn RecordReader>, RecorderError> {
    match format {
        RecordFormat::Protobuf => Ok(Box::new(ProtobufReader::new(path)?)),
        RecordFormat::Json => Ok(Box::new(JsonReader::new(path)?)),
    }
}

/// 创建默认的文件头部
pub fn create_default_header() -> RecordFileHeader {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos() as u64;
    
    RecordFileHeader {
        magic: RECORD_FILE_MAGIC,
        version: RECORD_FILE_VERSION.to_string(),
        created_timestamp_ns: now,
        system_context: Some(SystemContext {
            hostname: gethostname::gethostname().to_string_lossy().to_string(),
            process_id: std::process::id(),
            process_name: std::env::current_exe()
                .ok()
                .and_then(|p| p.file_name().map(|n| n.to_string_lossy().to_string()))
                .unwrap_or_else(|| "unknown".to_string()),
            cuda_version: "unknown".to_string(), // TODO: 获取实际CUDA版本
        }),
        statistics: Some(RecordStatistics::default()),
    }
}